from math import cos, sin, pi

def adjust_learning_rate(optimizer, current_epoch, max_epoch, lr_min=0., lr_max=0.01, warmup=True):
    warmup_epoch = 5 if warmup else 0
    lr = optimizer.param_groups[0]['lr']
    lr_min = lr * 0.5
    if current_epoch < warmup_epoch:
        lr = lr_max * current_epoch / warmup_epoch
    else:
        lr = lr_min + 0.5 * (lr_max - lr_min) * (
            1.0
            + cos(
                pi
                * (current_epoch - warmup_epoch)
                / (max_epoch - warmup_epoch)
            )
        )
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
        print("lr:", lr)
